import torch
import numpy as np
from model_wrapper.utils import convert_to_tensor
from sklearn.model_selection import train_test_split

if __name__ == "__main__":
    x1 = torch.tensor([0, 1, 2, 3, 4, 5, 6, 7, 8, 9])
    x2 = np.array([4, 5, 6, 7, 8, 9, 10, 11, 12, 13])
    y = torch.tensor([7, 8, 9, 10, 11, 12, 13, 14, 15, 16])
    X = (x1, x2)
    split_data = train_test_split(*X, y, test_size=0.2, random_state=42)
    print(split_data)
    y_train, y_test = split_data[-2:]
    # print(y_train, y_test)
    X = split_data[:-2]
    X_train, X_test = X[::2], X[1::2]
    # for i, x in enumerate(split_data[:-2]):
    #     if i % 2 == 0:
    #         X_train.append(x)
    #     else:
    #         X_test.append(x)
    
    print(tuple(X_train))
    print(tuple(X_test))

    # size = len(X[0])
    # X = tuple((x if torch.is_tensor(x) else convert_to_tensor(x)) for x in X)
    # print("X1:",X)
    # batch_size = 2
    # chunks = size // batch_size if size % batch_size == 0 else size // batch_size + 1

    # X = [torch.chunk(x, chunks, dim=0) for x in X]
    # print("X2:",X)
    # for x in zip(*X):
    #     print(x)
